from langchain.chat_models import ChatOpenAI
from PyCmpltrtok.common import sep

if '__main__' == __name__:
    
    chat_model = ChatOpenAI(
        streaming=True,
        verbose=True,
        callbacks=[],
        openai_api_key="token1",
        openai_api_base=f"http://127.0.0.1:6001/v1",
        model_name="chatglm2-6b",
        temperature=1.0,
        openai_proxy=None,
        top_p=1.0,
        max_tokens=2048,
    )

    from langchain.schema import HumanMessage

    text = "What would be a good company name for a company that makes colorful socks?"
    messages = [HumanMessage(content=text)]
    
    xresult = chat_model.invoke(messages)
    sep()
    print(xresult)
    sep()
    print(xresult.content)
    sep()
